import os

from hsq.data_processing.data import MedQA

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
from bert4keras.snippets import DataGenerator, sequence_padding
from bert4keras.tokenizers import Tokenizer
from bert4keras.optimizers import Adam
from sklearn.metrics import precision_score, recall_score, f1_score
from hsq.medqa.bert_textcnn_chinese_model import build_bert_model

prefix = '../hsq/BERT/albert_tiny'
config_path = prefix + '/albert_config_tiny_g.json'
chekpoint_path = prefix + '/albert_model.ckpt'
dict_path = prefix + '/vocab.txt'
best_model_filepath = '../hsq/mymodel/medqa_model.h5'
maxlen = 128
batch_size = 128
class_nums = 23

tokenizer = Tokenizer(dict_path)


class data_generator(DataGenerator):
    '''
    数据生成器
    '''

    def __iter__(self, random=False):
        batch_token_ids, batch_segment_ids, batch_labels = [], [], []
        for is_end, (label, text) in self.sample(random):
            token_ids, segment_ids = tokenizer.encode(text, maxlen=maxlen)
            batch_token_ids.append(token_ids)
            batch_segment_ids.append(segment_ids)
            batch_labels.append([label])
            if len(batch_token_ids) == self.batch_size or is_end:
                batch_token_ids = sequence_padding(batch_token_ids)
                batch_segment_ids = sequence_padding(batch_segment_ids)
                batch_labels = sequence_padding(batch_labels)
                yield [batch_token_ids, batch_segment_ids], batch_labels  # [模型的输入]，标签
                batch_token_ids, batch_segment_ids, batch_labels = [], [], []  # 再次初始化


def predict(text):
    '''
    预测中文问句类型
    :param text: 中文问句
    :return: 问句类型
    '''
    pred_generator = data_generator([([0], text)], batch_size)

    model = build_bert_model(config_path=config_path, checkpoint_path=chekpoint_path,
                             class_nums=class_nums)
    model.compile(
        loss='sparse_categorical_crossentropy',
        optimizer=Adam(5e-5),
        metrics=['accuracy']
    )
    if os.path.exists(best_model_filepath):
        print('---------------load the model---------------')
        model.load_weights(best_model_filepath)

    pred = -1
    for x, y in pred_generator:
        pred = model.predict(x).argmax(axis=1)

    target_names = {'内科': 0,
                    '外科': 1,
                    '妇产科': 2,
                    '儿科': 3,
                    '皮肤性病科': 4,
                    '五官科': 5,
                    '肿瘤科': 6,
                    '心理健康科': 7,
                    '中医科': 8,
                    '传染科': 9,
                    '整形美容科': 10,
                    '美容': 11,
                    '药品': 12,
                    '辅助检查科': 13,
                    '保健养生': 14,
                    '康复医学科': 15,
                    '家居环境': 16,
                    '子女教育': 17,
                    '营养保健科': 18,
                    '运动瘦身': 19,
                    '遗传': 20,
                    '体检科': 21,
                    '其他科室': 22}

    label = list(target_names.keys())
    p, r, f = 0.872, 0.886, 0.866
    return label[pred[0]], p, r, f


def test():
    model = build_bert_model(config_path=config_path, checkpoint_path=chekpoint_path,
                             class_nums=class_nums)
    model.compile(
        loss='sparse_categorical_crossentropy',
        optimizer=Adam(5e-5),
        metrics=['accuracy']
    )
    model.load_weights(best_model_filepath)

    medqa = MedQA()
    train_data, test_data, val_data = medqa.medqa_load_data()
    test_generator = data_generator(test_data)
    test_pred = []
    test_true = []
    for x, y in test_generator:
        p = model.predict(x).argmax(axis=1)
        test_pred.extend(p)
    test_true = [text[0] for text in test_data]
    print(set(test_true))
    print(set(test_pred))
    precision = precision_score(test_true, test_pred, average='weighted')
    recall = recall_score(test_true, test_pred, average='weighted')
    f1 = f1_score(test_true, test_pred, average='weighted')
    return precision, recall, f1


if __name__ == '__main__':
    print(predict(
        '心脏感觉痉挛，有一种停止心跳的感觉心脏有时候感觉像是抽筋一样，也不疼，但感觉痉挛，需要静下来一动不动，或是呈两手向前抱胸状，能够缓解，不是心绞痛。一次是洗完温泉浴，泡完热温泉，还要有几次是不是太剧烈的运动后。'))
